"""
===============================
@author     : yaoruiqi

@Time       : 2023/3/13 11:31

@version    : V

@introduce  :

@change     : 
===============================
"""
import socket
import struct
from base64 import b64decode
from hashlib import sha1

from Crypto.Cipher import AES
from loguru import logger


class MessageEncryptionAndDecryption:

    def __init__(self, token: str, encoding_aes_key: str):
        """
        消息加密解密初始化
        @param token: 微信开放平台上，服务方设置的接收消息的校验 token
        @param encoding_aes_key: 微信开放平台上配置的消息加解密密钥
        """
        self.token = token
        self.encoding_aes_key = encoding_aes_key

    def signature_verification(self, msg_timestamp: str, msg_nonce: str, msg_encrypt: str, msg_signature: str):
        """
        对消息体签名进行验证
        @param msg_timestamp: 请求URL上原有参数,时间戳
        @param msg_nonce: 请求URL上原有参数,随机数
        @param msg_encrypt: 密文消息体
        @param msg_signature: 微信返回的消息体签名
        """
        try:
            signature_text = ''.join(sorted([self.token, msg_timestamp, msg_nonce, msg_encrypt]))
            self_signature = sha1(signature_text.encode()).hexdigest()
            assert self_signature == msg_signature, '消息体签名验证错误'
            return True
        except AssertionError as e:
            logger.warning(f'{e}')
            return False
        except Exception as e:
            logger.exception(f'消息体签名验证失败，{e}')
            return False

    @staticmethod
    def decrypt_msg_encrypt(aes_key_decode: str, msg_encrypt: str, appid: str):
        """
        解密密文的消息体
        @param aes_key_decode: 32位的aes_key
        @param msg_encrypt: 密文消息体
        @param appid: appid，用于验证尾部的appid 是否正确
        @return:
        """
        try:
            cryptor = AES.new(aes_key_decode, AES.MODE_CBC, aes_key_decode[:16])
            plain_text = cryptor.decrypt(b64decode(msg_encrypt))
            content = plain_text[16:-plain_text[-1]]
            xml_len = socket.ntohl(struct.unpack("I", content[: 4])[0])
            xml_content = content[4: xml_len + 4].decode()
            from_appid = content[xml_len + 4:].decode()
            assert from_appid == appid, '消息体appid验证失败'
            return xml_content
        except AssertionError as _:
            return False
        except Exception as e:
            logger.exception(f'消息体解密失败，{e}')
            return False

    def decrypt(self, **kwargs):
        try:
            assert self.signature_verification(
                msg_timestamp=kwargs['msg_timestamp'],
                msg_nonce=kwargs['msg_nonce'],
                msg_encrypt=kwargs['msg_encrypt'],
                msg_signature=kwargs['msg_signature'],
            ), '消息签名验证失败'
            aes_key_decode = b64decode(self.encoding_aes_key + "=")
            assert len(aes_key_decode) == 32, 'aes_key长度不正确'
            xml_content = self.decrypt_msg_encrypt(aes_key_decode=aes_key_decode, msg_encrypt=kwargs['msg_encrypt'], appid=kwargs['appid'])
            assert xml_content, '消息体解密失败'
            return xml_content
        except Exception as e:
            logger.exception(f'{e}')
            return False
